// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <iosfwd>
#include <memory>
#include <string>
#include <vector>

#include <nebula/array/data.h>
#include <nebula/core/buffer.h>
#include <nebula/core/compare.h>

#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/bits/bit_util.h>
#include <turbo/base/macros.h>
#include <nebula/core/visitor.h>

namespace nebula {

    // ----------------------------------------------------------------------
    // User array accessor types

    /// \brief Array base type
    /// Immutable data array with some logical type and some length.
    ///
    /// Any memory is owned by the respective Buffer instance (or its parents).
    ///
    /// The base class is only required to have a null bitmap buffer if the null
    /// count is greater than 0
    ///
    /// If known, the null count can be provided in the base Array constructor. If
    /// the null count is not known, pass -1 to indicate that the null count is to
    /// be computed on the first call to null_count()
    class TURBO_EXPORT Array {
    public:
        virtual ~Array() = default;

        /// \brief Return true if value at index is null. Does not boundscheck
       [[nodiscard]] bool is_null(int64_t i) const { return !is_valid(i); }

        /// \brief Return true if value at index is valid (not null). Does not
        /// boundscheck
        [[nodiscard]] bool is_valid(int64_t i) const {
            if (null_bitmap_data_ != nullptr) {
                return bit_util::get_bit(null_bitmap_data_, i + data_->offset);
            }
            // Dispatching with a few conditionals like this makes is_null more
            // efficient for how it is used in practice. Making is_null virtual
            // would add a vtable lookup to every call and prevent inlining +
            // a potential inner-branch removal.
            if (type_id() == Type::SPARSE_UNION) {
                return !internal::is_null_sparse_union(*data_, i);
            }
            if (type_id() == Type::DENSE_UNION) {
                return !internal::is_null_dense_union(*data_, i);
            }
            if (type_id() == Type::RUN_END_ENCODED) {
                return !internal::is_null_run_end_encoded(*data_, i);
            }
            return data_->null_count != data_->length;
        }

        /// \brief Return a Scalar containing the value of this array at i
        [[nodiscard]] turbo::Result<std::shared_ptr<Scalar>> get_scalar(int64_t i) const;

        /// Size in the number of elements this array contains.
        [[nodiscard]] int64_t length() const { return data_->length; }

        /// A relative position into another array's data, to enable zero-copy
        /// slicing. This value defaults to zero
        [[nodiscard]] int64_t offset() const { return data_->offset; }

        /// The number of null entries in the array. If the null count was not known
        /// at time of construction (and set to a negative value), then the null
        /// count will be computed and cached on the first invocation of this
        /// function
        [[nodiscard]] int64_t null_count() const;

        /// \brief Computes the logical null count for arrays of all types including
        /// those that do not have a validity bitmap like union and run-end encoded
        /// arrays
        ///
        /// If the array has a validity bitmap, this function behaves the same as
        /// null_count(). For types that have no validity bitmap, this function will
        /// recompute the null count every time it is called.
        ///
        /// \see get_null_count
        [[nodiscard]] int64_t compute_logical_null_count() const;

        [[nodiscard]] const std::shared_ptr<DataType> &type() const { return data_->type; }

        [[nodiscard]] Type::type type_id() const { return data_->type->id(); }

        /// Buffer for the validity (null) bitmap, if any. Note that Union types
        /// never have a null bitmap.
        ///
        /// Note that for `null_count == 0` or for null type, this will be null.
        /// This buffer does not account for any slice offset
        [[nodiscard]] const std::shared_ptr<Buffer> &null_bitmap() const { return data_->buffers[0]; }

        /// Raw pointer to the null bitmap.
        ///
        /// Note that for `null_count == 0` or for null type, this will be null.
        /// This buffer does not account for any slice offset
        [[nodiscard]] const uint8_t *null_bitmap_data() const { return null_bitmap_data_; }

        /// Equality comparison with another array
        [[nodiscard]] bool equals(const Array &arr, const EqualOptions & = EqualOptions::defaults()) const;

        [[nodiscard]] bool equals(const std::shared_ptr<Array> &arr,
                    const EqualOptions & = EqualOptions::defaults()) const;

        /// \brief Return the formatted unified diff of nebula::Diff between this
        /// Array and another Array
        [[nodiscard]] std::string Diff(const Array &other) const;

        /// Approximate equality comparison with another array
        ///
        /// epsilon is only used if this is Fp32Array or Fp64Array
        [[nodiscard]] bool approx_equals(const std::shared_ptr<Array> &arr,
                          const EqualOptions & = EqualOptions::defaults()) const;

        [[nodiscard]] bool approx_equals(const Array &arr,
                          const EqualOptions & = EqualOptions::defaults()) const;

        /// Compare if the range of slots specified are equal for the given array and
        /// this array.  end_idx exclusive.  This methods does not bounds check.
        [[nodiscard]] bool range_equals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,
                         const Array &other,
                         const EqualOptions & = EqualOptions::defaults()) const;

        [[nodiscard]] bool range_equals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,
                         const std::shared_ptr<Array> &other,
                         const EqualOptions & = EqualOptions::defaults()) const;

        [[nodiscard]] bool range_equals(const Array &other, int64_t start_idx, int64_t end_idx,
                         int64_t other_start_idx,
                         const EqualOptions & = EqualOptions::defaults()) const;

        [[nodiscard]] bool range_equals(const std::shared_ptr<Array> &other, int64_t start_idx,
                         int64_t end_idx, int64_t other_start_idx,
                         const EqualOptions & = EqualOptions::defaults()) const;

        /// \brief Apply the ArrayVisitor::Visit() method specialized to the array type
        turbo::Status Accept(ArrayVisitor *visitor) const;

        /// Construct a zero-copy view of this array with the given type.
        ///
        /// This method checks if the types are layout-compatible.
        /// Nested types are traversed in depth-first order. Data buffers must have
        /// the same item sizes, even though the logical types may be different.
        /// An error is returned if the types are not layout-compatible.
        [[nodiscard]]  turbo::Result<std::shared_ptr<Array>> View(const std::shared_ptr<DataType> &type) const;

        /// \brief Construct a copy of the array with all buffers on destination
        /// Memory Manager
        ///
        /// This method recursively copies the array's buffers and those of its children
        /// onto the destination MemoryManager device and returns the new Array.
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> copy_to(const std::shared_ptr<MemoryManager> &to) const;

        /// \brief Construct a new array attempting to zero-copy view if possible.
        ///
        /// Like copy_to this method recursively goes through all of the array's buffers
        /// and those of it's children and first attempts to create zero-copy
        /// views on the destination MemoryManager device. If it can't, it falls back
        /// to performing a copy. See Buffer::ViewOrCopy.
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> view_or_copy_to(
                const std::shared_ptr<MemoryManager> &to) const;

        /// Construct a zero-copy slice of the array with the indicated offset and
        /// length
        ///
        /// \param[in] offset the position of the first element in the constructed
        /// slice
        /// \param[in] length the length of the slice. If there are not enough
        /// elements in the array, the length will be adjusted accordingly
        ///
        /// \return a new object wrapped in std::shared_ptr<Array>
        [[nodiscard]]  std::shared_ptr<Array> slice(int64_t offset, int64_t length) const;

        /// slice from offset until end of the array
        [[nodiscard]]  std::shared_ptr<Array> slice(int64_t offset) const;

        /// Input-checking variant of Array::slice
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> slice_safe(int64_t offset, int64_t length) const;

        /// Input-checking variant of Array::slice
        [[nodiscard]] turbo::Result<std::shared_ptr<Array>> slice_safe(int64_t offset) const;

        [[nodiscard]] const std::shared_ptr<ArrayData> &data() const { return data_; }

        [[nodiscard]] int num_fields() const { return static_cast<int>(data_->child_data.size()); }

        /// \return pretty_print representation of array suitable for debugging
        [[nodiscard]] std::string to_string() const;

        /// \brief Perform cheap validation checks to determine obvious inconsistencies
        /// within the array's internal data.
        ///
        /// This is O(k) where k is the number of descendents.
        ///
        /// \return turbo::Status
        turbo::Status validate() const;

        /// \brief Perform extensive validation checks to determine inconsistencies
        /// within the array's internal data.
        ///
        /// This is potentially O(k*n) where k is the number of descendents and n
        /// is the array length.
        ///
        /// \return turbo::Status
        turbo::Status validate_full() const;

        /// \brief Return the device_type that this array's data is allocated on
        ///
        /// This just delegates to calling device_type on the underlying ArrayData
        /// object which backs this Array.
        ///
        /// \return DeviceAllocationType
        [[nodiscard]]  DeviceAllocationType device_type() const { return data_->device_type(); }

    protected:
        Array() = default;
        TURBO_DEFAULT_MOVE_AND_ASSIGN(Array);

        std::shared_ptr<ArrayData> data_;
        const uint8_t *null_bitmap_data_ = nullptr;

        /// Protected method for constructors
        void set_data(const std::shared_ptr<ArrayData> &data) {
            if (data->buffers.size() > 0) {
                null_bitmap_data_ = data->GetValuesSafe<uint8_t>(0, /*offset=*/0);
            } else {
                null_bitmap_data_ = nullptr;
            }
            data_ = data;
        }

    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(Array);

        friend void print_to(const Array &x, std::ostream *os);
    };

    static inline std::ostream &operator<<(std::ostream &os, const Array &x) {
        os << x.to_string();
        return os;
    }

    /// Base class for non-nested arrays
    class TURBO_EXPORT FlatArray : public Array {
    protected:
        using Array::Array;
    };

    /// Base class for arrays of fixed-size logical types
    class TURBO_EXPORT PrimitiveArray : public FlatArray {
    public:
        PrimitiveArray(const std::shared_ptr<DataType> &type, int64_t length,
                       const std::shared_ptr<Buffer> &data,
                       const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                       int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        /// Does not account for any slice offset
        [[nodiscard]]  const std::shared_ptr<Buffer> &values() const { return data_->buffers[1]; }

    protected:
        PrimitiveArray() : raw_values_(nullptr) {}

        void set_data(const std::shared_ptr<ArrayData> &data) {
            this->Array::set_data(data);
            raw_values_ = data->GetValuesSafe<uint8_t>(1, /*offset=*/0);
        }

        explicit PrimitiveArray(const std::shared_ptr<ArrayData> &data) { set_data(data); }

        const uint8_t *raw_values_;
    };

    /// Degenerate null type Array
    class TURBO_EXPORT NullArray : public FlatArray {
    public:
        using TypeClass = NullType;

        explicit NullArray(const std::shared_ptr<ArrayData> &data) { set_data(data); }

        explicit NullArray(int64_t length);

    private:
        void set_data(const std::shared_ptr<ArrayData> &data) {
            null_bitmap_data_ = nullptr;
            data->null_count = data->length;
            data_ = data;
        }
    };

}  // namespace nebula
